import os
import numpy as np
import pandas as pd
from argparse import ArgumentParser
from shutil import copy, copyfile
from tqdm import tqdm

# modify here
dataset_path = "/home/xiangjianjian/dataset/gaofen/"

target_path = dataset_path

paths = ['A', 'B', 'label', 'list']


def copy():
    for i in tqdm(range(2000)):
        a_path = os.path.join(dataset_path, "images", f"{i+1}_1.png")
        a_path_target = os.path.join(target_path, "A", f"{i+1}.png")

        b_path = os.path.join(dataset_path, "images", f"{i+1}_2.png")
        b_path_target = os.path.join(target_path, "B", f"{i+1}.png")

        gt_path = os.path.join(dataset_path, "gt", f"{i+1}_change.png")
        gt_path_target = os.path.join(target_path, "label", f"{i+1}.png")

        try:
            copyfile(a_path, a_path_target)
            copyfile(b_path, b_path_target)
            copyfile(gt_path, gt_path_target)
        except Exception as e:
            print("Failed to copy, %s" % e)


def split_dataset(ratio=0.8):
    file_names = []
    for i in range(2000):
        file_names.append(f"{i+1}.png")
    file_names = np.array(file_names)
    np.random.shuffle(file_names)
    split_index = int(ratio * len(file_names))
    train_names = file_names[:split_index]
    val_names = file_names[split_index:]
    write_list(train_names, split="train")
    write_list(val_names, split="val")


def write_list(names, split="train"):
    file = open(
        os.path.join(target_path, "list", f"{split}.txt"), mode="w", encoding="utf-8"
    )
    for line in names:
        file.write(f"{line}\n")
    file.close()
    
def mkdirs():
    for path in paths:
        path = os.path.join(target_path, path)
        if not os.path.exists(path):
            os.mkdirs(path)


if __name__ == "__main__":
    mkdirs()
    copy()
    split_dataset(ratio=0.9)
